import json
import os
import sys
from functools import partial
from argparse import ArgumentParser

cwd = os.path.dirname(
    os.path.join(os.getcwd(), __file__)
)

sys.path.append(os.path.normpath(cwd+'/..'))

from tools.search_feature_points import get_crn
from tools.spn_helper import (
    merge_img2pt_chosen_pt, divide_img2pt, merged_img2pt_to_geojson, 
    merged_img2pt_to_win_gjson, chosen_func_nms
)

parser = ArgumentParser()
parser.add_argument(
    '-s', '--set_json', help='json of photo list'
)
parser.add_argument(
    '-i', '--img2pt_json', help='img2pt point list', 
    nargs='*'
)
parser.add_argument(
    '-m', '--merged_img2pt', help='merged img2pt point list', 
    nargs='*'
)
parser.add_argument(
    '-j', '--json_dir', help='dirname to save json'
)
parser.add_argument(
    '-g', '--gjson_dir', default='', help='dirname to save geojson'
)
parser.add_argument(
    '-p', '--prefix', default='', help='prefix of json and geojson'
)
parser.add_argument(
    '-c', '--cities', help='city name correspond to merged_img2pt', nargs='*'
)
parser.add_argument(
    '-r', '--rel_path', help='relative path correspond to merged_img2pt', 
    nargs='*'
)
parser.add_argument(
    '-n', '--mean', default=True, type=bool, 
    help='do we use mean of channels to find key points'
)
parser.add_argument(
    '-f', '--chosen_func', default='nms',
    help='function to choose keypoints from different modals'
)
args = parser.parse_args()

if args.set_json == None and args.img2pt_json == None and args.merged_img2pt == None:
    raise ValueError("all of set_json, img2pt_json and merged_img2pt are None!")

scheme = 'HARRIS'

# photo list -> img2pt
if args.img2pt_json == None and args.merged_img2pt == None:
    if args.set_json == None:
        raise ValueError("all file path are None!")
    with open(args.set_json) as fp:
        dset_dict = json.load(fp)
    
    if args.json_dir == None:
        raise ValueError("json_dir is needed!")

    modal_list = []
    for modal, crn_dict in dset_dict.items():
        if modal == 'PS-RGB':
            img_pt_dict = get_crn(crn_dict, 
                os.path.join(
                    args.json_dir, f"{args.prefix}_PS-RGB_{scheme}.json"
                ),
                type_='RGB', scheme=scheme, modal='PS-RGB', 
                border_radius=160, mean=args.mean)
            modal_list.append('PS-RGB')

        elif modal == 'SAR-Intensity':
            img_pt_dict = get_crn(crn_dict, 
                os.path.join(
                    args.json_dir, f"{args.prefix}_SAR_{scheme}.json"
                ),
                type_='SAR', scheme=scheme,
                border_radius=160, mean=args.mean)
            modal_list.append('SAR')

# merge img2pt
if args.merged_img2pt == None:
    if args.img2pt_json == None:
        try:
            img_pt_dict
        except:
            raise ValueError("No img2pt_json, no img2pt exists!")
        else:
            args.img2pt_json = []
            for modal in modal_list:
                args.img2pt_json.append(
                    os.path.join(args.json_dir, 
                        f"{args.prefix}_{modal}_{scheme}.json")
                )

    if args.json_dir == None:
        raise ValueError("json_dir is needed!")

    if args.chosen_func == 'nms':
        args.chosen_func = partial(chosen_func_nms, radius=160)
    elif args.chosen_func == 'None':
        args.chosen_func = None

    merge_img2pt_chosen_pt(args.img2pt_json, save_path=os.path.join(
        args.json_dir, f"{args.prefix}_{scheme}_merged.json"
    ), chosen_func=args.chosen_func)

# divide merged img2pt
    divide_img2pt(
        img2pt_json=os.path.join(
        args.json_dir, f"{args.prefix}_{scheme}_merged.json"
        ), save_path=os.path.join(
        args.json_dir, f"{args.prefix}_{scheme}_divided.json"
        )
    )

# merged img2pt -> geometry geojson
if args.merged_img2pt == None:
    args.merged_img2pt = [
        os.path.join(args.json_dir, f"{args.prefix}_{scheme}_merged.json")
    ]

if args.gjson_dir == None:
    raise ValueError("gjson_dir is needed!")

if args.cities == None:
    raise ValueError("cities is needed!")

if args.rel_path == None:
    args.rel_path = args.cities

merged_img2pt_to_geojson(
    img2pt_json=args.merged_img2pt, save_dir=os.path.join(
        args.gjson_dir, f"{args.prefix}_{scheme}_geometry.geojson"
    ), dset_dir=args.rel_path
)

# merged img2pt -> win geojson
merged_img2pt_to_win_gjson(
    args.cities, img2pt_json=args.merged_img2pt,
    save_dir=os.path.join(
        args.gjson_dir, f"{args.prefix}_{scheme}_windows.geojson"
    ),
    win_size=(320, 320), modals=['SAR-Intensity', 'PS-RGB'], 
)

print("-------------------finished-------------------")